package com.plugin.mybtaisxplus.util;

import com.intellij.database.model.DasColumn;
import com.intellij.database.model.DasTableKey;
import com.intellij.database.psi.DbTable;
import com.intellij.database.util.DasUtil;
import com.intellij.util.containers.JBIterable;
import com.plugin.mybtaisxplus.dialog.GeneratorSettings;
import org.apache.velocity.VelocityContext;

import java.text.SimpleDateFormat;
import java.util.*;

/**
 * 模板数据模型工厂
 * 负责从数据库元数据创建 Velocity 模板所需的数据上下文 (Context)。
 */
public class DataModelFactory {

    /**
     * 为 Entity 创建数据上下文。
     * @param table    数据库表对象
     * @param settings 用户配置
     * @param entityClassName 生成的实体类名
     * @return 包含所有 Entity 模板所需数据的 VelocityContext
     */
    public static VelocityContext createEntityContext(DbTable table, GeneratorSettings settings, String entityClassName) {
        String basePackage = settings.getBasePackage();
        String currentDate = new SimpleDateFormat("yyyy-MM-dd").format(new Date());

        List<Map<String, Object>> fields = new ArrayList<>();
        Set<String> imports = new HashSet<>();

        JBIterable<? extends DasColumn> columns = DasUtil.getColumns(table);
        for (DasColumn column : columns) {
            Map<String, Object> field = new HashMap<>();
            // 【修改】调用新的工具类
            String fieldName = GeneratorUtils.toCamelCase(column.getName(), false);
            String javaType = GeneratorUtils.mapJdbcTypeToFullName(column.getDataType().getSpecification());

            field.put("name", fieldName);
            field.put("type", javaType.substring(javaType.lastIndexOf(".") + 1)); // 获取简单类名
            field.put("remarks", column.getComment() == null ? "" : column.getComment());
            field.put("isStringType", "String".equals(field.get("type")));
            fields.add(field);

            // 【修改】对 byte[] 类型特殊处理，它不需要 import
            if (!javaType.startsWith("java.lang") && !javaType.equals("byte[]")) {
                imports.add(javaType);
            }
        }

        VelocityContext context = new VelocityContext();
        context.put("packageName", basePackage + "." + settings.getEntityPackage());
        context.put("imports", imports);
        context.put("date", currentDate);
        context.put("classInfo", Map.of("className", entityClassName, "remarks", table.getComment() == null ? "" : table.getComment()));
        context.put("fields", fields);
        return context;
    }

    /**
     * 为 DTO 创建数据上下文。
     * @param settings 用户配置
     * @param entityClassName 关联的实体类名
     * @param dtoClassName 生成的 DTO 类名
     * @param entityContext 从 Entity 上下文中复用部分数据
     * @return 包含所有 DTO 模板所需数据的 VelocityContext
     */
    @SuppressWarnings("unchecked")
    public static VelocityContext createDtoContext(GeneratorSettings settings, String entityClassName, String dtoClassName, VelocityContext entityContext) {
        String basePackage = settings.getBasePackage();
        Set<String> imports = new HashSet<>((Collection<String>) entityContext.get("imports"));
        imports.add("java.util.List");

        VelocityContext context = new VelocityContext();
        context.put("packageName", basePackage + "." + settings.getQueryPackage());
        context.put("imports", imports);
        context.put("date", entityContext.get("date"));
        context.put("entityClassName", entityClassName);
        context.put("dtoClassName", dtoClassName);
        context.put("fields", entityContext.get("fields"));
        return context;
    }
    /**
     * 【全面增强版】创建一个包含所有代码生成所需信息的全局数据模型。
     * 此方法是代码生成器的核心数据工厂，负责：
     * 1. 传递所有 UI 上的配置（包括注解、选项、模板集）。
     * 2. 动态处理 import 语句。
     * 3. 对字段进行精细化分析（主键、类型、驼峰命名等）。
     * 4. 创建动态排序所需的数据结构。
     *
     * @param table    当前处理的表
     * @param settings 用户的所有配置
     * @param i        当前是第几个表 (用于从settings中获取类名)
     * @return 一个包含所有模板变量的、信息完备的 VelocityContext
     */
    public static VelocityContext createGlobalContext(DbTable table, GeneratorSettings settings, int i) {
        VelocityContext context = new VelocityContext();
        String basePackage = settings.getBasePackage();
        String currentDate = new SimpleDateFormat("yyyy-MM-dd").format(new Date());

        // --- 1. 基础信息 ---
        context.put("basePackage", basePackage);
        context.put("date", currentDate);
        context.put("author", settings.getAuthor());
        context.put("tableName", table.getName());

        // --- 2. 各层类名/包名 ---
        String entityClassName = settings.getClassNames().get(i);
        String entityVarName = GeneratorUtils.toCamelCase(entityClassName, false);
        String dtoClassName = entityClassName + "QueryDTO";
        String mapperClassName = entityClassName + "Mapper";
        String serviceClassName = entityClassName + "Service";
        String serviceImplClassName = serviceClassName + "Impl";
        String controllerClassName = entityClassName + "Controller";

        context.put("entityClassName", entityClassName);
        context.put("entityVarName", entityVarName);
        context.put("dtoClassName", dtoClassName);
        context.put("mapperClassName", mapperClassName);
        context.put("serviceClassName", serviceClassName);
        context.put("serviceImplClassName", serviceImplClassName);
        context.put("controllerClassName", controllerClassName);
        context.put("serviceVarName", GeneratorUtils.toCamelCase(serviceClassName, false));
        context.put("requestMapping", "/" + entityVarName);

        context.put("entityPackage", basePackage + "." + settings.getEntityPackage());
        context.put("dtoPackage", basePackage + "." + settings.getQueryPackage());
        context.put("mapperPackage", basePackage + "." + settings.getMapperPackage());
        context.put("servicePackage", basePackage + "." + settings.getServicePackage());
        context.put("serviceImplPackage", basePackage + "." + settings.getServiceImplPackage());
        context.put("controllerPackage", basePackage + "." + settings.getControllerPackage());
        context.put("xmlPackage", settings.getXmlPackage());

        context.put("classInfo", Map.of(
                "remarks", table.getComment() == null ? "" : table.getComment(),
                "className", entityClassName
        ));

        // --- 3. Controller 父类与返回值包装 ---
        String superClass = settings.getSuperClass();
        if (superClass != null && !superClass.trim().isEmpty()) {
            context.put("superClass", superClass.trim());
            context.put("superClassName", superClass.substring(superClass.lastIndexOf('.') + 1));
        }
        String wrapperClass = settings.getResponseWrapperClass();
        if (wrapperClass != null && !wrapperClass.trim().isEmpty()) {
            context.put("isResponseWrapped", true);
            context.put("responseWrapperClass", wrapperClass.trim());
            context.put("responseWrapperSimpleName", wrapperClass.substring(wrapperClass.lastIndexOf('.') + 1));
            context.put("successMethodName", settings.getSuccessMethodName().trim());
            context.put("failureMethodName", settings.getFailureMethodName().trim());
        } else {
            context.put("isResponseWrapped", false);
        }

        // --- 4. 【核心】注解与选项配置 ---
        String annotationType = settings.getAnnotation();
        context.put("isMp2", GeneratorSettings.ANNOTATION_MP2.equals(annotationType));
        context.put("isMp3", GeneratorSettings.ANNOTATION_MP3.equals(annotationType));
        context.put("isJpa", GeneratorSettings.ANNOTATION_JPA.equals(annotationType));

        context.put("generateComments", settings.isGenerateComments());
        context.put("useActualColumnAnnotation", settings.isUseActualColumnAnnotation());
        context.put("generateToString", settings.isGenerateToString());
        context.put("useLombok", settings.isUseLombok());
        context.put("useJsr310", settings.isUseJsr310());

        // --- 5. 【核心】动态 Import 处理 ---
        Set<String> imports = new HashSet<>();
        if (settings.isUseLombok()) {
            imports.add("lombok.Data");
        }
        if (GeneratorSettings.ANNOTATION_MP2.equals(annotationType)) {
            imports.add("com.baomidou.mybatisplus.annotations.TableName");
            imports.add("com.baomidou.mybatisplus.annotations.TableId");
            imports.add("com.baomidou.mybatisplus.enums.IdType");
            if (settings.isUseActualColumnAnnotation()) {
                imports.add("com.baomidou.mybatisplus.annotations.TableField");
            }
        } else if (GeneratorSettings.ANNOTATION_MP3.equals(annotationType)) {
            imports.add("com.baomidou.mybatisplus.annotation.TableName");
            imports.add("com.baomidou.mybatisplus.annotation.TableId");
            imports.add("com.baomidou.mybatisplus.annotation.IdType");
            if (settings.isUseActualColumnAnnotation()) {
                imports.add("com.baomidou.mybatisplus.annotation.TableField");
            }
        } else if (GeneratorSettings.ANNOTATION_JPA.equals(annotationType)) {
            // 使用通配符导入，更简洁
            imports.add("javax.persistence.*");
        }

        // --- 6. 【核心】字段详细处理 ---
        List<Map<String, Object>> fields = new ArrayList<>();
        Map<String, String> javaFieldToDbColumnMap = new HashMap<>();
        DasTableKey primaryKey = DasUtil.getPrimaryKey(table);
        Set<String> pkColumnNames = new HashSet<>();
        if (primaryKey != null) {
            primaryKey.getColumnsRef().names().forEach(pkColumnNames::add);
        }

        Set<String> wrapperTypes = Set.of(
                "java.lang.Integer", "java.lang.Long", "java.lang.Double", "java.lang.Float",
                "java.math.BigDecimal", "java.lang.Boolean", "java.util.Date", "java.sql.Time"
        );

        JBIterable<? extends DasColumn> columns = DasUtil.getColumns(table);
        for (DasColumn column : columns) {
            Map<String, Object> field = new HashMap<>();
            String columnName = column.getName();
            String fieldName = GeneratorUtils.processColumnName(columnName, settings);
            String javaTypeFullName = GeneratorUtils.mapJdbcTypeToFullName(column.getDataType().getSpecification());
            String javaTypeSimpleName = javaTypeFullName.substring(javaTypeFullName.lastIndexOf('.') + 1);

            field.put("name", fieldName);
            field.put("type", javaTypeSimpleName);
            field.put("remarks", column.getComment() == null ? "" : column.getComment());
            field.put("columnName", columnName);
            field.put("isPrimaryKey", pkColumnNames.contains(columnName));

            // 精细化的类型判断，供模板使用
            boolean isString = "java.lang.String".equals(javaTypeFullName);
            field.put("isStringType", isString);
            field.put("isWrapperType", !isString && wrapperTypes.contains(javaTypeFullName));

            // 为 service.vm 中的动态排序逻辑准备数据
            javaFieldToDbColumnMap.put(fieldName, columnName);

            fields.add(field);

            // 添加字段类型自身的 import (如果需要)
            if (!javaTypeFullName.startsWith("java.lang") && !javaTypeFullName.equals("byte[]")) {
                imports.add(javaTypeFullName);
            }
        }
        context.put("fields", fields);
        context.put("imports", imports); // 包含所有动态添加的 import
        context.put("javaFieldToDbColumnMap", javaFieldToDbColumnMap);

        return context;
    }
}